import argparse
import random

def split_tsv(input_file, train_file, val_file, val_size):
    # 获取文件的总行数
    with open(input_file, 'r', encoding='utf-8') as file:
        total_lines = sum(1 for _ in file) 

    # 随机生成验证集的行号
    random.seed(42)
    val_indices = set(random.sample(range(total_lines), val_size))

    # 初始化计数器
    current_line = 0

    # 打开文件分别写入训练集和验证集
    with open(input_file, 'r', encoding='utf-8') as file, \
         open(train_file, 'w', encoding='utf-8') as train_out, \
         open(val_file, 'w', encoding='utf-8') as val_out:
        
        # 遍历每一行并分配到相应的文件
        for line in file:
            if current_line in val_indices:
                val_out.write(line)
            else:
                train_out.write(line)
            current_line += 1

def main():
    # 创建ArgumentParser对象
    parser = argparse.ArgumentParser(description="Split a TSV file into training and validation sets.")
    
    # 添加参数
    parser.add_argument("input_file", type=str, help="Path to the input TSV file")
    parser.add_argument("train_file", type=str, help="Path to the output training TSV file")
    parser.add_argument("val_file", type=str, help="Path to the output validation TSV file")
    parser.add_argument("val_size", type=int, help="Number of rows to use for the validation set")

    # 解析参数
    args = parser.parse_args()

    # 调用split_tsv函数
    split_tsv(args.input_file, args.train_file, args.val_file, args.val_size)

if __name__ == "__main__":
    main()
